import asyncio
import websockets
from threading import Thread
import threading
import cv2
from flask import Flask, request, jsonify
from flask_cors import CORS
import torch
from utils.datasets import LoadImages
from utils.general import non_max_suppression, scale_coords, check_img_size
from utils.torch_utils import select_device
from models.experimental import attempt_load
import os
from datetime import datetime

# 存储所有活动的 WebSocket 连接
active_connections = set()

# 创建一个全局的事件循环
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)

# 全局变量来存储模型
device = select_device('cpu')  # 选择设备类型，这里默认使用CPU
model = attempt_load(weights='mod/best.pt', map_location=device)  # 加载预训练模型
model_2 = attempt_load(weights='mod/best2.pt', map_location=device)
stride = int(model.stride.max())  # 获取模型处理图像时的最大步长
stride_2 = int(model_2.stride.max())
data_jg = {1: {}, 2: '未知', 3: 100}
time_data = datetime(2024, 12, 19, 12, 0, 0)
new_data = {
    '2a': 1, '2b': 2, '2c': 3, '2d': 4,
    '3a': 5, '3b': 6, '3c': 7, '3d': 8,
    '4a': 9, '4b': 10, '4c': 11, '4d': 12,
    '5a': 13, '5b': 14, '5c': 15, '5d': 16,
    '6a': 17, '6b': 18, '6c': 19, '6d': 20,
    '7a': 21, '7b': 22, '7c': 23, '7d': 24,
    '8a': 25, '8b': 26, '8c': 27, '8d': 28,
    '9a': 29, '9b': 30, '9c': 31, '9d': 32,
    '10a': 33, '10b': 34, '10c': 35, '10d': 36,
    '11a': 37, '11b': 38, '11c': 39, '11d': 40,
    '12a': 41, '12b': 42, '12c': 43, '12d': 44,
    '13a': 45, '13b': 46, '13c': 47, '13d': 48,
    '14a': 49, '14b': 50, '14c': 51, '14d': 52
}
set_data = {'state': 'yes'}
jl_data = {}


# 当客户端连接时调用
async def handler(websocket, path=None):
    active_connections.add(websocket)
    print(f"New connection: {id(websocket)}")
    try:
        async for message in websocket:
            print(f"Received message from {id(websocket)}: {message}")
            await websocket.send(f"Echo: {message}")
    finally:
        active_connections.remove(websocket)
        print(f"Connection {id(websocket)} closed.")


async def start_websocket():
    # 设置 WebSocket 服务器
    async with websockets.serve(handler, "0.0.0.0", 5000):
        await asyncio.Future()  # 阻塞，保持服务器运行


async def broadcast_data(message):
    print("message" + message);
    if active_connections:
        print(f"Broadcasting message to {len(active_connections)} clients: {message}")
        for connection in list(active_connections):
            try:
                await connection.send(message)
                print(f"Message sent to {id(connection)}")
            except websockets.ConnectionClosed:
                print(f"Connection closed: {id(connection)}")
                active_connections.remove(connection)
    else:
        print("No active WebSocket connections to broadcast message.")


# 创建 Flask 应用
app = Flask(__name__)
CORS(app)


@torch.no_grad()
def detect_image(source):
    for i in range(2):
        if i == 0:
            imgsz = check_img_size(640, s=stride_2)  # 对于 model 1 使用的步长
            dataset = LoadImages(source, img_size=imgsz, stride=stride_2, auto=True)  # 这里使用 model_2 的步长
            for path, img, im0s, _ in dataset:  # 遍历数据集中的图像
                img = torch.from_numpy(img).to(device)  # 将图像转换为 Tensor 并发送到设备
                img = img.half() if device.type != 'cpu' else img.float()  # 如果不是 CPU，则转换为半精度浮点数
                img = img / 255.0  # 归一化图像到 0.0-1.0
                if len(img.shape) == 3:
                    img = img[None]  # 如果图像没有批次维度，则添加一个批次维度

                # 使用 model_2 进行推理
                pred = model_2(img, augment=False)[0]  # 进行推理
                pred = non_max_suppression(pred, conf_thres=0.77, iou_thres=0.35, max_det=1000)  # 应用非极大值抑制
                for det in pred:  # 处理检测结果
                    if len(det):
                        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0s.shape).round()  # 将框坐标调整到原图大小
                        for *xyxy, conf, cls in reversed(det):
                            x1, y1, x2, y2 = map(int, xyxy)  # 将坐标转换为整数
                            cropped_img = im0s[y1:y2, x1:x2]  # 裁剪出目标区域
                            os.remove(source)
                            source = 'upload/new2.jpg'
                            cv2.imwrite(source, cropped_img)
                    else:
                        if os.path.exists(path):
                            os.remove(path)  # 删除处理过的文件
                        return
        elif i == 1:
            imgsz = check_img_size(640, s=stride)  # 对于 model 1 使用的步长
            dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=True)  # 这里使用 model 的步长

            for path, img, im0s, _ in dataset:  # 遍历数据集中的图像
                img = torch.from_numpy(img).to(device)  # 将图像转换为 Tensor 并发送到设备
                img = img.half() if device.type != 'cpu' else img.float()  # 如果不是 CPU，则转换为半精度浮点数
                img = img / 255.0  # 归一化图像到 0.0-1.0
                if len(img.shape) == 3:
                    img = img[None]  # 如果图像没有批次维度，则添加一个批次维度
                # 使用 model 进行推理
                pred = model(img, augment=False)[0]  # 进行推理
                pred = non_max_suppression(pred, conf_thres=0.75, iou_thres=0.35, max_det=1000)  # 应用非极大值
                if set_data['state'] != 'no':
                    for cropped_det in pred:
                        if len(cropped_det):
                            # if len(cropped_det) >= 11:

                            cropped_det[:, :4] = scale_coords(img.shape[2:], cropped_det[:, :4], im0s.shape).round()
                            for *xyxy, conf, cls in reversed(cropped_det):
                                names = model.module.names if hasattr(model, 'module') else model.names
                                label = f'{names[int(cls)]}'
                                if new_data[label] in jl_data:
                                    jl_data[new_data[label]] = jl_data[new_data[label]] + 1
                                else:
                                    jl_data[new_data[label]] = 1
                            if len(jl_data) >= 9:
                                data_jg[1] = [key for key, value in
                                              sorted(jl_data.items(), key=lambda item: item[1], reverse=True)[
                                              :13]]
                                data_jg[1].sort(reverse=True)
                            print(jl_data)


@app.route('/create/upload', methods=["GET", "POST"])
def upload():
    file = request.files.get("file")
    upload_folder = 'upload'

    if not os.path.exists(upload_folder):
        os.makedirs(upload_folder)

    file_path = os.path.join(upload_folder, file.filename)
    file.save(file_path)

    # 读取图像并调整分辨率（可选）
    img = cv2.imread(file_path)
    resized_img = cv2.resize(img, (720, 1080))
    output_path = '../upload/new1.jpg'
    cv2.imwrite(output_path, resized_img, [int(cv2.IMWRITE_JPEG_QUALITY), 20])

    # 启动图片处理线程
    thread = threading.Thread(target=detect_image, args=(output_path,))
    thread.start()

    # 图片上传完成后，立即通过 WebSocket 发送消息
    message = "Image uploaded successfully"
    asyncio.run_coroutine_threadsafe(broadcast_data(message), loop)

    return jsonify({"message": "Image uploaded successfully", "file_path": output_path})


# 在单独的线程中启动 Flask 应用
def start_flask():
    app.run(debug=False, host='0.0.0.0', port=5001)


# 主程序
if __name__ == "__main__":
    # 使用线程启动 Flask
    flask_thread = Thread(target=start_flask)
    flask_thread.start()

    # 启动 WebSocket 服务
    asyncio.run(start_websocket())
